
import os
import pymysql
from typing import List, Optional
import jinja2
from Config import DB_CONFIG, GENERATE_CONFIG
from FieldInfo import FieldInfo
from TableInfo import TableInfo

class MySQLGenerator:
    """MySQL代码生成器"""

    def __init__(self):
        self.connection = None
        self.tables: List[TableInfo] = []
        self.template_env = jinja2.Environment(
            loader=jinja2.FileSystemLoader('templates'),
            trim_blocks=True,
            lstrip_blocks=True
        )

    def connect(self) -> bool:
        """连接到MySQL数据库"""
        try:
            self.connection = pymysql.connect(
                host=DB_CONFIG['host'],
                port=DB_CONFIG['port'],
                user=DB_CONFIG['user'],
                password=DB_CONFIG['password'],
                database=DB_CONFIG['database'],
                charset=DB_CONFIG['charset'],
                cursorclass=pymysql.cursors.DictCursor
            )
            return True
        except Exception as e:
            print(f"连接数据库失败: {e}")
            return False

    def disconnect(self):
        """断开数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None

    def get_tables(self) -> List[str]:
        """获取数据库中的所有表"""
        if not self.connection:
            return []

        try:
            with self.connection.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                result = cursor.fetchall()

                # 提取表名
                tables = []
                for row in result:
                    table_name = list(row.values())[0]
                    tables.append(table_name)

                return tables
        except Exception as e:
            print(f"获取表列表失败: {e}")
            return []

    def get_table_info(self, table_name: str) -> Optional[TableInfo]:
        """获取表的结构信息"""
        if not self.connection:
            return None

        try:
            with self.connection.cursor() as cursor:
                # 获取表字段信息
                cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
                columns = cursor.fetchall()

                fields = []
                for col in columns:
                    field_name = col['Field']
                    type_name = col['Type']
                    is_nullable = col['Null'] == 'YES'
                    default_value = col['Default']
                    comment = col.get('Comment', '')

                    # 判断是否为主键
                    is_primary = col['Key'] == 'PRI'

                    field = FieldInfo(
                        name=field_name,
                        type_name=type_name,
                        is_primary=is_primary,
                        is_nullable=is_nullable,
                        default_value=default_value,
                        comment=comment
                    )
                    fields.append(field)

                return TableInfo(name=table_name, fields=fields)
        except Exception as e:
            print(f"获取表 {table_name} 的结构信息失败: {e}")
            return None

    def generate_classes(self):
        """生成所有表的映射类和Mapper类"""
        # 创建输出目录
        output_dir = GENERATE_CONFIG['output_dir']
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 获取所有表
        tables = self.get_tables()
        if not tables:
            print("数据库中没有表")
            return

        # 获取每个表的信息
        self.tables = []
        for table in tables:
            table_info = self.get_table_info(table)
            if table_info:
                self.tables.append(table_info)

        # 为每个表生成映射类和Mapper类
        for table in self.tables:
            self._generate_model_class(table)
            self._generate_mapper_class(table)

        # 生成__init__.py文件
        self._generate_init_file()

        print(f"成功生成 {len(self.tables)} 个表的映射类和Mapper类")

    def _generate_model_class(self, table: TableInfo):
        """生成映射类"""
        # 准备模板参数
        class_name = table.get_class_name()
        template_params = {
            'class_name': class_name,
            'model_class_name': class_name + 'Model',
            'mapper_class_name': class_name + 'Mapper',
            'table_name': table.name,
            'fields': table.fields,
            'base_module': 'BaseModel',
            'base_class': GENERATE_CONFIG['base_class']
        }

        # 加载模板
        template = self.template_env.get_template('ModelTemplate.py')

        # 渲染模板
        content = template.render(**template_params)

        # 写入文件
        output_dir = os.path.join(GENERATE_CONFIG['output_dir'], 'models')
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        file_path = os.path.join(output_dir, f"{class_name}.py")
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)

    def _generate_mapper_class(self, table: TableInfo):
        """生成Mapper类"""
        # 准备模板参数
        class_name = table.get_class_name()
        mapper_class_name = table.get_mapper_class_name()
        template_params = {
            'class_name': class_name,
            'model_class_name': class_name + 'Model',
            'mapper_class_name': class_name + 'Mapper',
            'table_name': table.name,
            'fields': table.fields
        }

        # 加载模板
        template = self.template_env.get_template('MapperTemplate.py')

        # 渲染模板
        content = template.render(**template_params)

        # 写入文件
        output_dir = os.path.join(GENERATE_CONFIG['output_dir'], 'mappers')
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        file_path = os.path.join(output_dir, f"{mapper_class_name}.py")
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)

    def _generate_init_file(self):
        """生成__init__.py文件"""
        # 生成models/__init__.py
        models_init_content = "# 模型包\n"
        for table in self.tables:
            class_name = table.get_class_name()
            models_init_content += f"from .{class_name} import {class_name}\n"

        models_init_path = os.path.join(GENERATE_CONFIG['output_dir'], 'models', '__init__.py')
        with open(models_init_path, 'w', encoding='utf-8') as f:
            f.write(models_init_content)

        # 生成mappers/__init__.py
        mappers_init_content = "# Mapper包\n"
        for table in self.tables:
            mapper_class_name = table.get_mapper_class_name()
            mappers_init_content += f"from .{mapper_class_name} import {mapper_class_name}\n"

        mappers_init_path = os.path.join(GENERATE_CONFIG['output_dir'], 'mappers', '__init__.py')
        with open(mappers_init_path, 'w', encoding='utf-8') as f:
            f.write(mappers_init_content)

        # 生成根目录的__init__.py
        root_init_content = """# MySQL 生成器生成的代码

from .models import *
from .mappers import *

# 创建数据库连接
def create_connection():
    import pymysql
    return pymysql.connect(
        host='{}',
        port={},
        user='{}',
        password='{}',
        database='{}',
        charset='{}',
        cursorclass=pymysql.cursors.DictCursor
    )
""".format(
    DB_CONFIG['host'],
    DB_CONFIG['port'],
    DB_CONFIG['user'],
    DB_CONFIG['password'],
    DB_CONFIG['database'],
    DB_CONFIG['charset']
)

        root_init_path = os.path.join(GENERATE_CONFIG['output_dir'], '__init__.py')
        with open(root_init_path, 'w', encoding='utf-8') as f:
            f.write(root_init_content)
